import time
import torch
import model
import config
import torch.nn as nn
import torch.optim as optim
import DataSet
from torchvision.utils import save_image


# 默认配置
opt = config.DefaultConfig()

# 网络加载
NetG = model.generator(128)
NetD = model.discriminator(128)
NetG.weight_init(mean=0.0, std=0.002)
NetD.weight_init(mean=0.0, std=0.002)
NetG.cuda()
NetD.cuda()
#数据集
tarin_data = DataSet.train_loader

# 损失函数
BCE_loss = nn.BCELoss()

# 优化器
G_optimizer = optim.Adam(NetG.parameters(), lr=opt.lr, betas=(0.5, 0.999))
D_optimizer = optim.Adam(NetD.parameters(), lr=opt.lr, betas=(0.5, 0.999))

# 标签
onehot = torch.zeros(10, 10)
onehot = onehot.scatter_(1, torch.LongTensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]).view(10,1), 1).view(10, 10, 1, 1)
fill = torch.zeros([10, 10, opt.img_size, opt.img_size])
for i in range(10):
    fill[i, i, :, :] = 1

# 真假标签
y_real, y_fake = torch.ones(opt.batch_size).cuda(), torch.zeros(opt.batch_size).cuda()
# 固定标签
fix_noise = torch.randn(80, 100, 1, 1)
fix_label = torch.tensor([i for i in range(10) for j in range(8)])
fix_label = onehot[fix_label]

print("开始训练")
start_time = time.time()
for epoch in range(opt.train_epoch):
    for ii, (img, label) in enumerate(tarin_data):
        # 训练鉴别器  真图片
        NetD.zero_grad()
        real_img = img.cuda()
        out_put = NetD(real_img, fill[label].cuda())
        D_real_loss = BCE_loss(out_put.view(-1), y_real)
        # 训练鉴别器  假图片
        z_ = torch.randn(opt.batch_size, 100, 1, 1).cuda()
        y_ = torch.randint(0, 10, (opt.batch_size, 1))      # 生成标签
        y_label = onehot[y_.view(-1)].cuda()
        y_fill = fill[y_.view(-1)].cuda()

        fake_img = NetG(z_, y_label)
        fake_out = NetD(fake_img, y_fill)
        D_fake_loss = BCE_loss(fake_out.view(-1), y_fake)
        D_train_loss = D_real_loss + D_fake_loss

        D_train_loss.backward()
        D_optimizer.step()                  # 更新梯度

        # 训练生成器
        NetG.zero_grad()
        z_ = torch.randn(opt.batch_size, 100, 1, 1).cuda()
        y_ = torch.randint(0, 10, (opt.batch_size, 1))  # 生成标签
        y_label = onehot[y_.view(-1)].cuda()
        y_fill = fill[y_.view(-1)].cuda()

        img = NetG(z_, y_label)
        out = NetD(img, y_fill)

        G_train_loss = BCE_loss(out.view(-1), y_real)
        G_train_loss.backward()
        G_optimizer.step()
        if (ii+1) % 10 == 0:
            print(str(ii) + ":     GLoss:  " + str(G_train_loss.data) + "DLoss:  " + str(D_train_loss.data))

    NetG.eval()
    imgs = NetG(fix_noise.cuda(), fix_label.cuda())
    NetG.train()
    save_image(imgs, "./img/epoch"+str(epoch)+".png")
